#!/usr/local/bin/python3
# -*- coding: utf-8 -*-


"""
2022-11-06 by Dr Jie Zheng
有些IDL程序的格式实在是不规则，或者说，原作者的规则我觉得看不下去，所以写这么一个来进行格式化
基于常用的方式来处理，包括识别注释和代码，运算符两段空格，代码层次缩进等
"""


import re
import sys
import datetime


def idlfmt(infile):
    """
    2022-11-06 by Dr Jie Zheng
    有些IDL程序的格式实在是不规则，或者说，原作者的规则我觉得看不下去，所以写这么一个来进行格式化
    基于常用的方式来处理，包括识别注释和代码，运算符两段空格，代码层次缩进等
    """
    # 读文件
    lines = open(infile).readlines()
    
    # 初始化层次等
    indent = 0
    linejoin = False
    
    # 逐行处理
    for li in lines:
        # 先处理注释，可能的bug：字符串中的分号
        lis = li.split(";")
        # 命令部分
        com = lis[0].strip()
        # 注释部分，注意可能有多个分号，都要留下
        rem = ";".join(lis[1:]).strip()
        
        # 跳过字符串部分，利用引号拆分，统一用单引号
        comsp = com.replace('\"', '\'').split('\'')
        
        # 1、3、5等部分是字符串内容，不处理，0、2、4、6等是代码
        for i in range(0, len(comsp), 2):
            c = comsp[i]
            # 代码空格处理
            # 处理 +,- 避开指数浮点数，处理 ## 避免和 # 混淆
            c = re.sub(r"([\d\.][eEdD])\+", lambda m:m.group(1)+"➕", c)
            c = re.sub(r"([\d\.][eEdD])\-", lambda m:m.group(1)+"➖", c)
            c = re.sub(r"##", "￥", c)
            # 处理符号左边空格
            c = re.sub(r"(^|[^\s])([=\+\-\*/\^\<\>#￥])", lambda m:m.group(1)+" "+m.group(2), c)
            # 处理符号右边空格，比左边的多一个逗号
            c = re.sub(r"([=\+\-\*/\^\<\>#×\,])([^\s]|$)", lambda m:m.group(1)+" "+m.group(2), c)
            # 恢复指数浮点数以及##
            c = c.replace("➕", "+").replace("➖", "-").replace("￥", "##")
            comsp[i] = c
        com = '\''.join(comsp)
                
        # 开头 end 减少一层缩进
        if com.upper().startswith("END"):
            indent -= 1
        
        # 如果结尾的$前面没有空格，加一个
        com = re.sub(r"([^\s])\$$", lambda m:m.group(1)+" $", com)
        
        # 最后，剔除两侧的不必要的空格
        com = com.strip()
        
        # 根据当前缩进层次继续输出
        print(
            "  " * indent + 
            com + 
            ("  " if com and rem else "") +
            ("; " if rem else "") +
            rem
        )
        
        # 续行
        if com.endswith("$"):
            # 如果已经在续行模式，那么就不增加缩进
            if not linejoin:
                indent += 1
                linejoin = True
            else:
                pass
        else:
            # 如果是从续行转回普通，要处理，否则不管
            if linejoin:
                indent -= 1
                linejoin = False
            else:
                pass
        
        # 结尾是begin的，下一行增加一级缩进
        if com.upper().endswith("BEGIN"):
            indent += 1
        
        # pro 和 function 下一行开始缩进
        if com.upper().startswith("PRO ") or com.upper().startswith("FUNCTION "):
            indent += 1
    
    strnow = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S")
    print(f";;;;;  {strnow}  {infile}")


if __name__ == "__main__":
    for f in sys.argv[1:]:
        idlfmt(f)
